import torch
import time
import os

params_dict = {}
params_dict['device'] = torch.device("cpu")
# params_dict['device'] = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
params_dict['batch_size'] = 16
params_dict['nw'] = 0
params_dict['train_img_path'] = "./data/images/imgs"
params_dict['train_file_path'] = "./data/bladder_train.json"
params_dict['valid_img_path'] = "./data/images/imgs"
params_dict['valid_file_path'] = "./data/bladder_val.json"
params_dict['test_img_path'] = "./data/images/imgs"
params_dict['test_file_path'] = "./data/bladder_test.json"
params_dict['vocab_size'] = 1000
params_dict['embedding_dim'] = 64
params_dict['n_head'] = 4
params_dict['num_layers'] = 2
params_dict['v_vector_size'] = 64
params_dict['input_dim'] = 511
params_dict['hidden_dim'] = 511
params_dict['output_dim'] = 4  # 4分类
params_dict['is_train'] = True  # 是否训练
params_dict['epoch'] = 500
params_dict['lr'] = 0.0001
params_dict['pre_train'] = False
if params_dict['pre_train'] == True:
    params_dict['pre_train_model_path'] = ""
params_dict['save_file_path'] = './run/PNAT' + str(time.strftime("%Y-%m-%d %H_%M_%S", time.localtime()))
params_dict['save_best_model_path'] = os.path.join(params_dict['save_file_path'], "best.pt")
params_dict['save_last_model_path'] = os.path.join(params_dict['save_file_path'], "last.pt")

